package mybatisPlugin.config;

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLObject;
import com.alibaba.druid.sql.ast.expr.SQLBetweenExpr;
import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlInsertStatement;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUpdateStatement;
import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitorAdapter;
import com.alibaba.druid.sql.visitor.SQLEvalVisitorUtils;
import com.alibaba.druid.util.JdbcConstants;
import mybatisPlugin.entity.SqlInfo;
import org.springframework.util.StringUtils;

import java.util.List;

import static mybatisPlugin.config.SqlAesUtils.aesEncryptString;
import static mybatisPlugin.config.SqlAesUtils.isProperty;

/**
 * CustomizeSqlASTVisitor
 * 加密字段作为参数时暂认定为其不会被函数处理
 *
 * @author 方典典
 * @date 2019/1/21 17:23
 */
public class CustomizeSqlASTVisitor extends MySqlASTVisitorAdapter {

    /**
     * 是否已经访问过where语句
     */
    private boolean hasVisitWhere;
    /**
     * 是否经过了加密处理的标记名称
     */
    public static final String ENCRYPT_FLAG = "_encrypt";
    /**
     * sql信息
     */
    private SqlInfo sqlInfo;

    public CustomizeSqlASTVisitor() {
        super();
        sqlInfo = new SqlInfo();
    }

    @Override
    public boolean visit(SQLBetweenExpr x) {
        SQLExpr testExpr = x.getTestExpr();
        if (!isProperty(testExpr)) {
            return super.visit(x);
        }
        String key = SQLUtils.toMySqlString(testExpr);
        if (!key.endsWith(ENCRYPT_FLAG)) {
            return super.visit(x);
        }
        x.setBeginExpr(aesEncryptString(x.getBeginExpr()));
        x.setEndExpr(aesEncryptString(x.getEndExpr()));
        sqlInfo.setNeedAesHandle(true);
        return super.visit(x);
    }

    @Override
    public boolean visit(SQLBinaryOpExpr x) {
        SQLObject sqlObject = x.getParent();
        boolean noParam = isProperty(x.getLeft()) && isProperty(x.getRight());
        boolean notProperty = !isProperty(x.getLeft()) && !isProperty(x.getRight());
        if (x.getOperator().getPriority() != 110 || notProperty || noParam) {
            return super.visit(x);
        }
        String value = SQLUtils.toMySqlString(x.getLeft());
        if (!StringUtils.isEmpty(value) && value.endsWith(ENCRYPT_FLAG)) {
            x.setRight(aesEncryptString(x.getRight()));
            sqlInfo.setNeedAesHandle(true);
            while (sqlObject.getParent() != null) {
                sqlObject = sqlObject.getParent();
            }
            if (sqlObject instanceof MySqlUpdateStatement) {
                hasVisitWhere = true;
                this.visit((MySqlUpdateStatement) sqlObject);
            }
            sqlInfo.setSql(SQLUtils.toMySqlString(sqlObject));
        }
        return super.visit(x);
    }

    @Override
    public boolean visit(SQLExprTableSource x) {
        String alia = x.getAlias();
        String table = x.getName().getSimpleName();
        if (!StringUtils.isEmpty(alia)) {
            table = table + " " + alia;
        }
        sqlInfo.getTables().add(table);
        return true;
    }

    @Override
    public boolean visit(MySqlSelectQueryBlock x) {
        List<SQLSelectItem> selectList = x.getSelectList();
        SelectSqlPropertyASTVisitor selectSqlPropertyASTVisitor = new SelectSqlPropertyASTVisitor();
        selectList.forEach(sqlSelectItem -> {
            sqlSelectItem.accept(selectSqlPropertyASTVisitor);
            sqlInfo.setNeedAesHandle(selectSqlPropertyASTVisitor.isNeedAesHandle());
        });
        sqlInfo.setSql(SQLUtils.toMySqlString(x));
        return super.visit(x);
    }

    @Override
    public boolean visit(MySqlUpdateStatement x) {
        sqlInfo.setWhere(SQLUtils.toMySqlString(x.getWhere()));
        if (hasVisitWhere) {
            return true;
        }
        if(x.getTableSource() instanceof SQLJoinTableSource){
            sqlInfo.setComplex(true);
        }
        List<SQLUpdateSetItem> itemList = x.getItems();
        itemList.forEach(item -> {
            String key = SQLUtils.toMySqlString(item.getColumn());
            if (key.endsWith(ENCRYPT_FLAG)) {
                item.setValue(aesEncryptString((item.getValue())));
                sqlInfo.setNeedAesHandle(true);
            }
            String value;
            try {
                value = SQLEvalVisitorUtils.evalExpr(JdbcConstants.MYSQL, SQLUtils.toMySqlString(item.getValue()))
                        .toString();
            } catch (Exception e) {
                //解析失败
                sqlInfo.setComplex(true);
                value = SQLUtils.toMySqlString(item.getValue());
            }
            sqlInfo.getItems().put(key, value);
        });
        sqlInfo.setSql(SQLUtils.toMySqlString(x));
        return true;
    }

    @Override
    public boolean visit(MySqlInsertStatement x) {
        insertItemsHandle(x.getColumns(), x.getValuesList());
        sqlInfo.setSql(SQLUtils.toMySqlString(x));
        if (!x.getDuplicateKeyUpdate().isEmpty()) {
            sqlInfo.setReplace(true);
            sqlInfo.setComplex(true);
        }
        return true;
    }

    @Override
    public boolean visit(SQLReplaceStatement x) {
        insertItemsHandle(x.getColumns(), x.getValuesList());
        sqlInfo.setSql(SQLUtils.toMySqlString(x));
        sqlInfo.setReplace(true);
        sqlInfo.setComplex(true);
        return true;
    }

    /**
     * insertItemsHandle
     *
     * @param columns
     * @param valuesClauseList
     * @return void
     * @author 方典典
     * @date 2019/1/21 17:05
     */
    private void insertItemsHandle(List<SQLExpr> columns, List<SQLInsertStatement.ValuesClause> valuesClauseList) {
        for (int i = 0; i < columns.size(); i++) {
            String column = SQLUtils.toMySqlString(columns.get(i));
            if (column.endsWith(ENCRYPT_FLAG)) {
                for (SQLInsertStatement.ValuesClause valuesClause : valuesClauseList) {
                    valuesClause.getValues().set(i, aesEncryptString(valuesClause.getValues().get(i)));
                    sqlInfo.setNeedAesHandle(true);
                }
            }
        }
    }

    public SqlInfo getSqlInfo() {
        return sqlInfo;
    }

}